# cogview_image_server.py

import os
import json
import base64
import asyncio
from typing import Optional
from dotenv import load_dotenv
from mcp.server import FastMCP
from zhipuai import ZhipuAI

app = FastMCP('cogview_image_server')

# 初始化智谱AI客户端
# 请设置环境变量 ZHIPUAI_API_KEY 或直接在这里替换为你的API密钥
load_dotenv() 
API_KEY = os.getenv('ZHIPUAI_API_KEY', 'your-api-key-here')
client = ZhipuAI(api_key=API_KEY)
print(API_KEY)


@app.tool()
async def generate_image_with_cogview(
    prompt: str,
    size: str = "1024x1024",
    quality: str = "standard",
    n: int = 1
) -> str:
    """
    使用 CogView-3-Flash 生成图片
    
    Args:
        prompt: 图片描述提示词，支持中英文
        size: 图片尺寸，支持 "1024x1024", "768x1024", "1024x768" 等
        quality: 图片质量，"standard" 或 "hd"
        n: 生成图片数量，默认为1
    
    Returns:
        生成的图片信息JSON字符串
    """
    try:
        # 调用智谱AI的CogView-3-Flash模型
        response = client.images.generations(
            model="cogview-3-flash",  # 使用CogView-3-Flash模型
            prompt=prompt,
            size=size,
            quality=quality,
            n=n
        )
        
        # 处理响应
        if response and response.data:
            result = {
                "success": True,
                "images": []
            }
            
            for image_data in response.data:
                image_info = {
                    "url": getattr(image_data, 'url', None),
                    "b64_json": getattr(image_data, 'b64_json', None),
                    "revised_prompt": getattr(image_data, 'revised_prompt', None)
                }
                result["images"].append(image_info)
            
            return json.dumps(result, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                "success": False,
                "error": "未收到有效响应"
            }, ensure_ascii=False)
            
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": f"生成图片失败: {str(e)}"
        }, ensure_ascii=False)


@app.tool()
async def generate_image_advanced_cogview(
    prompt: str,
    model: str = "cogview-3-flash",
    size: str = "1024x1024",
    quality: str = "standard",
    style: Optional[str] = None,
    user_id: Optional[str] = None
) -> str:
    """
    使用 CogView 高级参数生成图片
    
    Args:
        prompt: 图片描述提示词
        model: 模型名称，可选 "cogview-3-flash", "cogview-3-plus" 等
        size: 图片尺寸
        quality: 图片质量
        style: 图片风格（如果支持）
        user_id: 用户ID（用于追踪）
    
    Returns:
        生成的图片信息JSON字符串
    """
    try:
        # 构建请求参数
        params = {
            "model": model,
            "prompt": prompt,
            "size": size,
            "quality": quality
        }
        
        # 添加可选参数
        if style:
            params["style"] = style
        if user_id:
            params["user"] = user_id
        
        # 调用API
        response = client.images.generations(**params)
        
        if response and response.data:
            result = {
                "success": True,
                "model": model,
                "images": []
            }
            
            for image_data in response.data:
                image_info = {
                    "url": getattr(image_data, 'url', None),
                    "b64_json": getattr(image_data, 'b64_json', None),
                    "revised_prompt": getattr(image_data, 'revised_prompt', None)
                }
                result["images"].append(image_info)
            
            return json.dumps(result, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                "success": False,
                "error": "未收到有效响应"
            }, ensure_ascii=False)
            
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": f"生成图片失败: {str(e)}"
        }, ensure_ascii=False)


@app.tool()
async def save_image_from_url(
    image_url: str,
    filename: Optional[str] = None
) -> str:
    """
    从URL下载并保存图片到本地
    
    Args:
        image_url: 图片URL
        filename: 保存的文件名，如果不提供则自动生成
    
    Returns:
        保存结果信息
    """
    try:
        import httpx
        from datetime import datetime
        
        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"cogview_image_{timestamp}.png"
        
        async with httpx.AsyncClient() as http_client:
            response = await http_client.get(image_url)
            response.raise_for_status()
            
            # 确保目录存在
            os.makedirs("generated_images", exist_ok=True)
            filepath = os.path.join("generated_images", filename)
            
            with open(filepath, 'wb') as f:
                f.write(response.content)
            
            return json.dumps({
                "success": True,
                "filepath": filepath,
                "filename": filename,
                "size": len(response.content)
            }, ensure_ascii=False)
            
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": f"保存图片失败: {str(e)}"
        }, ensure_ascii=False)


@app.tool()
async def save_image_from_base64(
    b64_data: str,
    filename: Optional[str] = None
) -> str:
    """
    从base64数据保存图片到本地
    
    Args:
        b64_data: base64编码的图片数据
        filename: 保存的文件名
    
    Returns:
        保存结果信息
    """
    try:
        from datetime import datetime
        
        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"cogview_image_{timestamp}.png"
        
        # 解码base64数据
        image_data = base64.b64decode(b64_data)
        
        # 确保目录存在
        os.makedirs("generated_images", exist_ok=True)
        filepath = os.path.join("generated_images", filename)
        
        with open(filepath, 'wb') as f:
            f.write(image_data)
        
        return json.dumps({
            "success": True,
            "filepath": filepath,
            "filename": filename,
            "size": len(image_data)
        }, ensure_ascii=False)
        
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": f"保存图片失败: {str(e)}"
        }, ensure_ascii=False)


@app.tool()
async def batch_generate_images(
    prompts: list,
    size: str = "1024x1024",
    quality: str = "standard"
) -> str:
    """
    批量生成图片
    
    Args:
        prompts: 提示词列表
        size: 图片尺寸
        quality: 图片质量
    
    Returns:
        批量生成结果
    """
    try:
        results = []
        
        for i, prompt in enumerate(prompts):
            try:
                response = client.images.generations(
                    model="cogview-3-flash",
                    prompt=prompt,
                    size=size,
                    quality=quality,
                    n=1
                )
                
                if response and response.data:
                    image_data = response.data[0]
                    result = {
                        "index": i,
                        "prompt": prompt,
                        "success": True,
                        "url": getattr(image_data, 'url', None),
                        "b64_json": getattr(image_data, 'b64_json', None)
                    }
                else:
                    result = {
                        "index": i,
                        "prompt": prompt,
                        "success": False,
                        "error": "未收到有效响应"
                    }
                
                results.append(result)
                
                # 添加延迟以避免API限制
                await asyncio.sleep(1)
                
            except Exception as e:
                results.append({
                    "index": i,
                    "prompt": prompt,
                    "success": False,
                    "error": str(e)
                })
        
        return json.dumps({
            "success": True,
            "total": len(prompts),
            "results": results
        }, ensure_ascii=False, indent=2)
        
    except Exception as e:
        return json.dumps({
            "success": False,
            "error": f"批量生成失败: {str(e)}"
        }, ensure_ascii=False)


if __name__ == '__main__':
    # 检查API密钥
    if API_KEY == 'your-api-key-here':
        print("警告: 请设置 ZHIPUAI_API_KEY 环境变量或修改代码中的API密钥")
        print("使用方法:")
        print("export ZHIPUAI_API_KEY='your-actual-api-key'")
        print("python cogview_image_server.py")
    
    app.run(transport='stdio')